# This is a mep2d, intended for creating 2-dimensional potential 
# energy surfaces of molecular systems.
# Copyright (C) 2011  Frank Uhlig
# 
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

#!/usr/bin/env python

from matplotlib import rc
import matplotlib

rc('text', usetex=True)
rc('font', family='serif')
rc('font', size=12)

# figsize = (8, 6)
# dpi = 300

import numpy as np
import os,sys
import constants

#from fes import plotstuff
from plotter import plot3d
from grid_data import cgrid

import matplotlib.pyplot as mpl
import pylab as p
import enthought.mayavi.mlab as M
import mpl_toolkits.mplot3d as p3
from mpl_toolkits.mplot3d import axes3d

def plot3dd(X,Y,Z,plotkind,xlabel,ylabel,zlabel,levels):
    """plot free energy surface in different ways"""

    if plotkind.lower().strip() == 'wire':
        fig = mpl.figure()
        ax = axes3d.Axes3D(fig)
        #ax = fig.add_subplot(111,projection='3d')
        ax.plot_wireframe(X,Y,Z)
        mpl.show()
        #mpl.savefig('image.png')
    elif plotkind.lower().strip() == 'mesh3d':
        M.mesh(X,Y,Z)
        M.show()
        #M.savefig('image.png')
    elif plotkind.lower().strip() == 'contour':
        p.contour(X,Y,Z,levels)
        p.figure()
        p.contourf(X,Y,Z,levels)
        p.show()
        #p.savefig('image.png')
    elif plotkind.lower().strip() == 'contour3d':
        fig=p.figure()
        ax = p3.Axes3D(fig)
        ax.contourf3D(X,Y,Z)
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)
        ax.set_zlabel(xlabel)
        fig.add_axes(ax)
        p.show()
        #p.savefig('image.png')
    elif plotkind.lower().strip() == 'jet':
        p.hold(True)
        # Creating image
        p.imshow(Z, origin='lower', extent=[X.min(),X.max(),Y.min(),Y.max()])
        # Plotting contour lines
        p.contour(Z, origin='lower', extent=[X.min(),X.max(),Y.min(),Y.max()])
        p.xlabel(xlabel)
        p.ylabel(ylabel)
        p.title(zlabel)
        # Adding a line plot slicing the z matrix just for fun. 
        #plot(X[:], Z[50, :])
        p.show()
        #p.savefig()
        

if len(sys.argv) >1:
    plotkind = sys.argv[1]
else:
    plotkind = 'wire'

allen = np.array([],dtype='float')
minen = np.array([],dtype='float')

#xlabel='$\mathrm{H_{2}O \: O-H \: [\AA]}$'
#ylabel='$\mathrm{H_{3}O \: O-H \: [\AA]}$'
xlabel='$\mathrm{H_{2}O \: [\AA]}$'
ylabel='$\mathrm{H_{3}O \: [\AA]}$'
zlabel='Total energy'

start = 10
stop = 22
step = 1

X,Y,Z = cgrid(start,stop,step,start,stop,step)

i = -1
for dist in range(start,stop,step): #<-- H3O-OH
    i += 1
    j = -1
    for distt in range(start,stop,step): #<-- H2O-OH
        j += 1
        f = open('../calcs/' + str(dist) + '-w-' + str(distt) + '/gradient','r') 
        for line in f.readlines():
            if line.strip().startswith('cycle'):
                #print line
                energy = float(line.split()[6])
                #print energy
        if plotkind == 'mesh3d':
            Z[i,j] = energy*100
        else:
            Z[i,j] = energy
            #print Z[i,:],i,j
        f.close()

Z[:,:] = (Z[:,:] - Z[1,1])*constants.h2ev
#Z[:,:] = (Z[:,:] - np.min(Z[:,:]))*constants.h2ev

#print np.min(Z),np.max(Z),np.arange(np.min(Z),np.max(Z),(np.max(Z)-np.min(Z))/100.)
#plot3dd(X,Y,Z,plotkind,xlabel,ylabel,zlabel,np.arange(np.min(Z),np.max(Z),(np.max(Z)-np.min(Z)/50.)))

# X *= 1/10.
# Y *= 1/10.
plot3d(X,Y,Z,plotkind,xlabel,ylabel,zlabel)

#allowed values for plotkind are: wire, mesh3d, contour, contour3d, jet; default is wire
